import cv2
import numpy as np
import os
import time
import threading
from concurrent.futures import ThreadPoolExecutor
import atexit
import datetime
from collections import deque

class MergedPosePositionVisualizer:
    """高效的合并Pose和Position可视化器 - 专为实时性能优化"""

    def __init__(self,rect_size=(0.8, 2.5), scale_factor=100, enable_display=True):
        """
        初始化合并可视化器
        
        Args:
            rect_size: 矩形大小（宽, 长），单位为米，默认为(0.8, 2.5)
            scale_factor: 米到像素的转换因子，默认为100（1米=100像素）
            enable_display: 是否启用显示功能，默认为True
        """
        # Position部分的参数
        self.rect_size = rect_size
        self.scale_factor = scale_factor
        self.rect_width, self.rect_height = rect_size
        
        # 画布配置 - 使用固定尺寸避免缩放
        self.scene_width = 1280  # 左边场景区域：固定1280x720
        self.scene_height = 720
        self.position_width = 400  # 右边位置区域：400像素宽
        self.canvas_width = self.scene_width + self.position_width  # 总宽度1680
        self.canvas_height = self.scene_height  # 总高度720
        
        # 显示控制
        self.enable_display = enable_display
        self.window_created = False
        self.window_name = "Merged Visualization"
        
        # 预分配画布 - 关键优化
        self.canvas = np.zeros((self.canvas_height, self.canvas_width, 3), dtype=np.uint8)
        self.scene_canvas = self.canvas[:, :self.scene_width]  # 1280x720区域
        self.position_canvas = self.canvas[:, self.scene_width:]  # 400x720区域
        
        # 位置图的中心点和比例
        self.pos_center_x = self.position_width // 2
        self.pos_center_y = self.scene_height // 2
        self.pos_scale = min(self.position_width, self.scene_height) // 4
        
        # Pose部分的颜色定义
        self.pose_colors = {
            'left_top': (0, 255, 0),    # 绿色
            'left_0': (255, 0, 255),    # 紫色
            'left_bottom': (255, 0, 0),  # 蓝色
            'bottom_0': (0, 128, 128),   # 深青色
            'right_bottom': (0, 255, 255), # 黄色
            'right_0': (255, 255, 0),    # 青色
            'right_top': (0, 0, 255),    # 红色
            'top_0': (128, 0, 128),      # 深紫色
        }
        
        # 整体颜色配置
        self.colors = {
            'standard_pose': (0, 255, 0),      # 绿色
            'current_pose': (0, 0, 255),       # 红色  
            'standard_rect': (0, 255, 255),    # 黄色
            'current_rect': (255, 0, 255),     # 紫色
            'background': (40, 40, 40),        # 深灰背景
            'grid': (80, 80, 80),              # 网格线
            'text': (255, 255, 255),           # 白色文字
            'split_line': (128, 128, 128)      # 分割线
        }
        
        # 定义连接顺序
        self.connect_order = [
            'left_top', 'left_0', 'left_bottom',
            'bottom_0', 'right_bottom', 'right_0',
            'right_top', 'top_0', 'left_top'
        ]
        
        # 字体设置
        self.font = cv2.FONT_HERSHEY_SIMPLEX
        self.font_scale = 0.4
        self.font_thickness = 1
        
        # 输出目录
        self.output_dir = "results_merged_visualization"
        os.makedirs(self.output_dir, exist_ok=True)
        
        # 性能监控
        self.frame_times = deque(maxlen=30)
        self.fps = 0
        
        # 图像缓存优化
        self._image_cache = {}
        self._size_cache = {}
        self.cache_max_size = 5  # 减少缓存大小，因为背景图在变化
        
        # 预计算的矩形点缓存
        self._rect_cache = {}
        
        # 预绘制静态元素
        self._prepare_static_elements()
        
        # 注册清理函数
        atexit.register(self.cleanup)
        
        print(f"高效合并可视化器初始化完成，输出目录: {self.output_dir}")

    def _prepare_static_elements(self):
        """预绘制静态元素（网格、标题等）"""
        # 清空画布为深色背景
        self.canvas[:] = self.colors['background']
        
        # 绘制分割线
        cv2.line(self.canvas, (self.scene_width, 0), (self.scene_width, self.scene_height), 
                self.colors['split_line'], 2)
        
        # 绘制标题
        cv2.putText(self.canvas, "POSE", (20, 30), cv2.FONT_HERSHEY_SIMPLEX, 
                   1, self.colors['text'], 2)
        cv2.putText(self.canvas, "POSITION", (self.scene_width + 20, 30), 
                   cv2.FONT_HERSHEY_SIMPLEX, 1, self.colors['text'], 2)
        
        # 在位置区域绘制网格
        self._draw_position_grid()
        
        # 保存静态背景
        self.static_background = self.canvas.copy()

    def _draw_position_grid(self):
        """绘制位置区域的网格"""
        # 绘制十字参考线
        cv2.line(self.position_canvas, 
                (self.pos_center_x, 0), (self.pos_center_x, self.scene_height),
                self.colors['grid'], 1)
        cv2.line(self.position_canvas,
                (0, self.pos_center_y), (self.position_width, self.pos_center_y), 
                self.colors['grid'], 1)
        
        # 绘制同心圆作为距离参考
        for radius in [50, 100, 150]:
            if radius < min(self.pos_center_x, self.pos_center_y):
                cv2.circle(self.position_canvas, 
                          (self.pos_center_x, self.pos_center_y), 
                          radius, self.colors['grid'], 1)

    def _create_display_window(self):
        """创建显示窗口（只创建一次）"""
        if not self.enable_display or self.window_created:
            return
            
        try:
            cv2.namedWindow(self.window_name, cv2.WINDOW_NORMAL)
            cv2.resizeWindow(self.window_name, self.canvas_width, self.canvas_height)
            self.window_created = True
        except Exception as e:
            print(f"无法创建显示窗口: {e}")
            self.enable_display = False

    def _process_background_image(self, image_input):
        """超高效处理背景图像 - 无缩放版本"""
        if image_input is None:
            return None
            
        # 处理图像输入
        if isinstance(image_input, (str, os.PathLike)):
            image_path = str(image_input)
            
            # 检查缓存
            if image_path in self._image_cache:
                image = self._image_cache[image_path]
            else:
                image = cv2.imread(image_path)
                if image is None:
                    return None
                
                # 有限缓存
                if len(self._image_cache) < self.cache_max_size:
                    self._image_cache[image_path] = image
                    
        elif isinstance(image_input, np.ndarray):
            image = image_input
        else:
            return None
        
        # 检查尺寸是否匹配
        h, w = image.shape[:2]
        if h == self.scene_height and w == self.scene_width:
            # 完美匹配，直接返回
            return image
        elif h == self.scene_height and w != self.scene_width:
            # 高度匹配，宽度裁剪或填充
            if w > self.scene_width:
                # 裁剪中间部分
                start_x = (w - self.scene_width) // 2
                return image[:, start_x:start_x + self.scene_width]
            else:
                # 宽度不足，居中填充
                result = np.zeros((self.scene_height, self.scene_width, 3), dtype=np.uint8)
                start_x = (self.scene_width - w) // 2
                result[:, start_x:start_x + w] = image
                return result
        else:
            # 尺寸不匹配，只在必要时才resize
            return cv2.resize(image, (self.scene_width, self.scene_height))

    def _draw_background_scene(self, processed_image):
        """绘制场景背景 - 直接拷贝版本"""
        if processed_image is not None:
            # 直接拷贝到场景区域（已经是正确尺寸）
            self.scene_canvas[:] = processed_image

    def _draw_pose_points(self, pose_dict, base_color, is_reference=False):
        """高效绘制pose点和连线"""
        if pose_dict is None:
            return
        
        # 收集有效点
        valid_points = []
        
        for name in self.connect_order[:-1]:  # 去掉重复的最后一个点
            if name in pose_dict and pose_dict[name]:
                x, y = pose_dict[name]
                point = (int(x), int(y))
                valid_points.append(point)
                
                # 绘制点
                if is_reference:
                    cv2.circle(self.scene_canvas, point, 3, (192, 192, 192), -1)
                else:
                    # 使用单独颜色
                    color = self.pose_colors.get(name, base_color)
                    cv2.circle(self.scene_canvas, point, 5, color, -1)
                    
                    # 添加标签（简化版，只有关键点）
                    if name in ['left_top', 'right_top', 'left_bottom', 'right_bottom']:
                        cv2.putText(self.scene_canvas, name.replace('_', ''), 
                                  (point[0]+8, point[1]), self.font, 
                                  self.font_scale, color, self.font_thickness)
        
        # 绘制连线（闭合多边形）
        if len(valid_points) >= 3:
            cv2.polylines(self.scene_canvas, [np.array(valid_points, dtype=np.int32)], 
                         True, base_color, 2 if is_reference else 3)

    def _compute_rectangle_points(self, vector):
        """快速计算矩形四个角点"""
        if vector is None or 'center' not in vector:
            return None
            
        center_x = vector['center'][0] * self.scale_factor + self.pos_center_x
        center_y = vector['center'][1] * self.scale_factor + self.pos_center_y
        
        # 默认矩形大小（像素）
        rect_w = self.rect_size[0] * self.scale_factor
        rect_h = self.rect_size[1] * self.scale_factor
        
        # 简化处理方向向量
        if 'left_vec' in vector and 'bottom_vec' in vector:
            left_vec = np.array(vector['left_vec'][:2])
            bottom_vec = np.array(vector['bottom_vec'][:2])
            
            # 快速标准化
            left_len = np.linalg.norm(left_vec)
            if left_len > 1e-6:
                left_vec = left_vec / left_len * rect_h
            else:
                left_vec = np.array([0, rect_h])
                
            bottom_len = np.linalg.norm(bottom_vec)
            if bottom_len > 1e-6:
                bottom_vec = bottom_vec / bottom_len * rect_w
            else:
                bottom_vec = np.array([rect_w, 0])
        else:
            # 默认方向
            left_vec = np.array([0, rect_h])
            bottom_vec = np.array([rect_w, 0])
        
        # 计算四个角点
        center = np.array([center_x, center_y])
        bottom_left = center
        bottom_right = center + bottom_vec
        top_left = center + left_vec  
        top_right = top_left + bottom_vec
        
        return np.array([bottom_left, bottom_right, top_right, top_left], dtype=np.int32)

    def _draw_rectangle(self, vector, color, thickness=2):
        """在位置区域绘制矩形"""
        points = self._compute_rectangle_points(vector)
        if points is not None:
            cv2.polylines(self.position_canvas, [points], True, color, thickness)

    def visualize_merged(self, pose_dict, st_pose_dict, standard_vector, test_vector, 
                        image_input, timestamp=None, display=True, save_async=True):
        """
        高效的合并可视化函数 - 优化为实时性能
        
        Args:
            pose_dict: 当前pose字典
            st_pose_dict: 标准pose字典
            standard_vector: 标准位置向量
            test_vector: 测试位置向量
            image_input: 图像路径(str)或图像数组(numpy.ndarray)
            timestamp: 时间戳，用作文件名(可选)
            display: 是否显示图像（默认为True）
            save_async: 是否异步保存（默认为True）
        """
        frame_start = time.time()
        
        try:
            # 1. 快速恢复静态背景
            self.canvas[:] = self.static_background
            
            # 2. 高效处理背景图像
            processed_image = self._process_background_image(image_input)
            
            # 3. 绘制场景背景
            if processed_image is not None:
                self._draw_background_scene(processed_image)
            
            # 4. 绘制pose（标准位置先画，当前位置后画以突出显示）
            if st_pose_dict:
                self._draw_pose_points(st_pose_dict, self.colors['standard_pose'], is_reference=True)
            
            if pose_dict:
                self._draw_pose_points(pose_dict, self.colors['current_pose'], is_reference=False)
            
            # 5. 绘制位置矩形
            if standard_vector:
                self._draw_rectangle(standard_vector, self.colors['standard_rect'], 2)
                
            if test_vector:
                self._draw_rectangle(test_vector, self.colors['current_rect'], 3)
            
            # 6. 计算并显示性能信息
            frame_time = time.time() - frame_start
            self.frame_times.append(frame_time)
            
            if len(self.frame_times) > 1:
                avg_frame_time = sum(self.frame_times) / len(self.frame_times)
                self.fps = 1.0 / max(avg_frame_time, 1e-6)
            
            # 显示性能信息
            fps_text = f"FPS: {self.fps:.1f} | {frame_time*1000:.1f}ms"
            cv2.putText(self.canvas, fps_text, (20, self.scene_height - 20), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.6, self.colors['text'], 1)
            
            # 7. 显示图像
            if display and self.enable_display:
                self._create_display_window()
                if self.window_created:
                    try:
                        cv2.imshow(self.window_name, self.canvas)
                        key = cv2.waitKey(1) & 0xFF
                        if key == ord('q'):
                            print("用户请求退出")
                    except Exception as e:
                        print(f"显示图像失败: {e}")
            
            # 8. 保存图像（如果需要）
            output_path = None
            if save_async and timestamp:
                filename = f"merged_{timestamp}.png"
                output_path = os.path.join(self.output_dir, filename)
                
                # 快速保存（降低质量以提高速度）
                cv2.imwrite(output_path, self.canvas, [cv2.IMWRITE_JPEG_QUALITY, 70])
            
            return output_path, self.canvas.copy()
            
        except Exception as e:
            print(f"可视化过程中发生错误: {e}")
            return None, None

    def batch_visualize(self, data_list, max_workers=2):
        """
        批量处理可视化（减少并发数以避免资源竞争）
        """
        results = []
        total = len(data_list)
        
        print(f"开始批量处理 {total} 个图像...")
        
        for i, data in enumerate(data_list, 1):
            result, _ = self.visualize_merged(
                data.get('pose_dict'),
                data.get('st_pose_dict'),
                data.get('standard_vector'),
                data.get('test_vector'),
                data.get('image_input'),
                data.get('timestamp'),
                display=False,  # 批量处理时不显示
                save_async=False
            )
            
            if result:
                results.append(result)
            
            if i % 10 == 0:
                print(f"处理进度: {i}/{total}")
        
        print(f"批量处理完成，成功处理 {len(results)} 个图像")
        return results

    def clear_cache(self):
        """清空缓存"""
        self._image_cache.clear()
        self._size_cache.clear()
        self._rect_cache.clear()

    def cleanup(self):
        """清理资源"""
        try:
            if self.window_created:
                cv2.destroyWindow(self.window_name)
            
            self.clear_cache()
        except Exception as e:
            print(f"清理资源时出错: {e}")

    def __del__(self):
        """析构函数"""
        self.cleanup()


# 保持向后兼容性的别名
class OptimizedPositionVisualizer(MergedPosePositionVisualizer):
    """向后兼容的PositionVisualizer"""
    def visualize_position(self, standard_vector, test_vector, image_input, timestamp=None):
        result, _ = self.visualize_merged(None, None, standard_vector, test_vector, image_input, timestamp)
        return result

class OptimizedPoseVisualizer(MergedPosePositionVisualizer):
    """向后兼容的PoseVisualizer"""
    def visualize_pose(self, pose_dict, st_pose_dict, image_input, output_dir=None, timestamp=None):
        if output_dir and output_dir != self.output_dir:
            self.output_dir = output_dir
            os.makedirs(self.output_dir, exist_ok=True)
        result, _ = self.visualize_merged(pose_dict, st_pose_dict, None, None, image_input, timestamp)
        return result

class PositionVisualizer(OptimizedPositionVisualizer):
    """向后兼容的PositionVisualizer"""
    pass

class PoseVisualizer(OptimizedPoseVisualizer):
    """向后兼容的PoseVisualizer"""
    pass


# 性能测试和使用示例
def quick_performance_test():
    """快速性能测试"""
    print("=== 快速性能测试 ===")
    
    # 创建测试数据
    test_vector = {
        'center': [1.0, 1.5],
        'left_vec': [0, 2.5, 0],
        'bottom_vec': [0.8, 0, 0]
    }
    
    standard_vector = {
        'center': [0.5, 1.0],
        'left_vec': [0, 2.5, 0],
        'bottom_vec': [0.8, 0, 0]
    }
    
    pose_dict = {
        'left_top': (100, 100),
        'left_0': (100, 200),
        'left_bottom': (100, 300),
        'bottom_0': (200, 300),
        'right_bottom': (300, 300),
        'right_0': (300, 200),
        'right_top': (300, 100),
        'top_0': (200, 100)
    }
    
    st_pose_dict = {
        'left_top': (120, 120),
        'left_0': (120, 220),
        'left_bottom': (120, 320),
        'bottom_0': (220, 320),
        'right_bottom': (320, 320),
        'right_0': (320, 220),
        'right_top': (320, 120),
        'top_0': (220, 120)
    }
    
    # 创建测试图像（固定1280x720尺寸）
    test_images = []
    for i in range(5):
        img = np.random.randint(0, 255, (720, 1280, 3), dtype=np.uint8)
        test_images.append(img)
    
    # 创建可视化器
    visualizer = MergedPosePositionVisualizer(enable_display=True)
    
    # 预热
    print("预热中...")
    for i in range(3):
        visualizer.visualize_merged(
            pose_dict, st_pose_dict, 
            standard_vector, test_vector, 
            test_images[i % len(test_images)],
            display=False
        )
    
    # 性能测试
    print("开始性能测试...")
    times = []
    
    for i in range(50):  # 测试50帧
        start_time = time.time()
        
        # 模拟变化的数据
        angle = i * 0.1
        offset_x = int(20 * np.sin(angle))
        offset_y = int(15 * np.cos(angle))
        
        # 变化的pose
        current_pose = {}
        for key, (x, y) in pose_dict.items():
            current_pose[key] = (x + offset_x, y + offset_y)
        
        # 变化的vector
        current_vector = {
            'center': [test_vector['center'][0] + 0.2 * np.sin(angle), 
                      test_vector['center'][1] + 0.1 * np.cos(angle)],
            'left_vec': test_vector['left_vec'],
            'bottom_vec': test_vector['bottom_vec']
        }
        
        # 调用可视化
        visualizer.visualize_merged(
            current_pose, st_pose_dict, 
            standard_vector, current_vector, 
            test_images[i % len(test_images)],
            timestamp=f"test_{i:03d}",
            display=(i < 10),  # 只显示前10帧
            save_async=False
        )
        
        elapsed = time.time() - start_time
        times.append(elapsed)
    
    # 统计结果
    times = np.array(times)
    avg_time = np.mean(times) * 1000
    min_time = np.min(times) * 1000
    max_time = np.max(times) * 1000
    
    theoretical_fps = 1000 / avg_time if avg_time > 0 else 0
    
    print(f"\n性能测试结果 (50帧):")
    print(f"平均帧时间: {avg_time:.2f}ms")
    print(f"最快帧时间: {min_time:.2f}ms")
    print(f"最慢帧时间: {max_time:.2f}ms")
    print(f"理论最大FPS: {theoretical_fps:.1f}")
    
    if avg_time < 16.7:
        print("✅ 性能优秀！可以支持60FPS实时显示")
    elif avg_time < 33.3:
        print("✅ 性能良好！可以支持30FPS流畅显示")
    else:
        print("⚠️ 性能一般，可能需要进一步优化")
    
    visualizer.cleanup()
    return avg_time


def realtime_demo():
    """实时演示"""
    print("\n=== 实时演示 ===")
    print("按 'q' 键退出演示")
    
    # 创建可视化器
    visualizer = MergedPosePositionVisualizer(enable_display=True)
    
    # 基础数据
    st_pose_dict = {
        'left_top': (150, 100), 'left_0': (150, 200), 
        'left_bottom': (150, 300), 'bottom_0': (250, 300),
        'right_bottom': (350, 300), 'right_0': (350, 200),
        'right_top': (350, 100), 'top_0': (250, 100)
    }
    
    standard_vector = {
        'center': [0, 0],
        'left_vec': [0, 2.5, 0],
        'bottom_vec': [0.8, 0, 0]
    }
    
    print("开始实时演示...")
    start_time = time.time()
    frame_count = 0
    
    try:
        while True:
            # 生成动画数据
            t = time.time() - start_time
            
            # 变化的背景图（固定1280x720）
            bg = np.random.randint(50, 200, (720, 1280, 3), dtype=np.uint8)
            
            # 动画pose
            offset_x = int(30 * np.sin(t * 2))
            offset_y = int(20 * np.cos(t * 1.5))
            
            pose_dict = {}
            for key, (x, y) in st_pose_dict.items():
                pose_dict[key] = (x + offset_x, y + offset_y)
            
            # 动画vector
            test_vector = {
                'center': [0.5 * np.sin(t), 0.3 * np.cos(t * 1.2)],
                'left_vec': [np.sin(t * 0.5) * 2.5, np.cos(t * 0.5) * 2.5, 0],
                'bottom_vec': [np.cos(t * 0.3) * 0.8, np.sin(t * 0.3) * 0.8, 0]
            }
            
            # 更新可视化
            result, frame = visualizer.visualize_merged(
                pose_dict, st_pose_dict,
                standard_vector, test_vector,
                bg,
                display=True,
                save_async=False
            )
            
            frame_count += 1
            
            # 检查退出
            key = cv2.waitKey(1) & 0xFF
            if key == ord('q'):
                break
                
            # 限制演示时间
            if t > 30:  # 30秒后自动结束
                print("演示时间结束")
                break
                
    except KeyboardInterrupt:
        print("\n用户中断演示")
    
    elapsed = time.time() - start_time
    avg_fps = frame_count / elapsed if elapsed > 0 else 0
    
    print(f"\n演示统计:")
    print(f"总时间: {elapsed:.1f}s")
    print(f"总帧数: {frame_count}")
    print(f"平均FPS: {avg_fps:.1f}")
    
    visualizer.cleanup()


if __name__ == "__main__":
    # 快速性能测试
    avg_time = quick_performance_test()
    
    # 如果性能足够好，运行实时演示
    if avg_time < 50:
        realtime_demo()
    else:
        print("\n性能不足以运行流畅的实时演示")